# Author : ZZH
# Date : 2025/5/26
from ultralytics import YOLO
from pathlib import Path
import argparse


def args_parser():
    parser = argparse.ArgumentParser(description="基于YOLO的安全帽推理脚本")

    parser.add_argument("--weights", type = str, default = "train6-20250523-175204-yolo11m-best.pt", help = "模型权重路径")
    parser.add_argument("--source", type = str, default = "0", help = "输入源（图像/文件夹/视频/摄像头ID，如‘0’）")
    parser.add_argument("--imgsz", type = int, default = 640, help = "输入图像尺寸")
    parser.add_argument("--conf", type = float, default = 0.25, help = "置信度阈值")
    parser.add_argument("--iou", type = float, default = 0.45, help = "IOU阈值")
    parser.add_argument("--save", type = bool, default = True, help = "保存预测图像")
    parser.add_argument("--save_dir", type = str, default = "runs/val", help = "保存预测图像的目录")
    parser.add_argument("--save_txt", type = bool, default = True, help = "保存预测结果为TXT")
    parser.add_argument("--save_conf", type = bool, default = True, help = "在TXT中包含置信度值")
    parser.add_argument("--save_frames", type = bool, default = True, help = "保存摄像头/视频每帧图像")
    parser.add_argument("--save_crop", type = bool, default = True, help = "保存检测框裁剪图像")
    parser.add_argument("--show", type = bool, default = True, help = "显示结果")

    return parser.parse_args()


def main():
    args = args_parser()

    model_path = Path(args.weights)
    if not model_path.is_absolute():
        model_path = Path(r"../models/checkpoints/") / args.weights

    source = args.source
    if not source.isdigit():
        source_path = Path(source)
        # 检查输入源路径是否存在
        if not source_path.exists():
            raise FileNotFoundError(f"输入源不存在：{source_path}")
        source = str(source_path)
    # ========== 模型验证 ==========
    if not model_path.exists():
        raise FileNotFoundError(f"模型文件不存在：{model_path}")
    # 加载模型
    model = YOLO(str(model_path))

    results = model.predict(
        source=source,
        imgsz=args.imgsz,
        conf=args.conf,
        show=args.show,
        iou=args.iou,
        save=args.save,
        save_txt=args.save_txt,
        save_conf=args.save_conf,
        save_crop=args.save_crop,
        save_frames=args.save_frames,
        project="runs/predict",
        name="exp"
    )

    print(f"推理完成，结果已保存至：{results.save_dir}")


if __name__ == "__main__":
    main()
